// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/container/f14_map.h>
#include <pollux/expression/complex_view_types.h>
#include <pollux/functions/udf.h>

namespace kumo::pollux::functions {
    template<typename TExec>
    struct MultimapFromEntriesFunction {
        POLLUX_DEFINE_FUNCTION_TYPES(TExec);

        MELON_ALWAYS_INLINE void call(
            out_type<Map<Generic<T1>, Array<Generic<T2> > > > &out,
            const arg_type<Array<Row<Generic<T1>, Generic<T2> > > > &inputArray) {
            // Reuse map and vector between rows to avoid re-allocating memory. The
            // benchmark shows 20-30% performance improvement.
            keyValuesMap_.clear();

            uniqueKeys_.clear();
            uniqueKeys_.reserve(inputArray.size());

            for (const auto &entry: inputArray) {
                POLLUX_USER_CHECK(entry.has_value(), "map entry cannot be null");
                const auto &key = entry.value().template at<0>();
                const auto &value = entry.value().template at<1>();

                POLLUX_USER_CHECK(key.has_value(), "map key cannot be null");

                auto result = keyValuesMap_.insert({key.value(), {}});
                result.first->second.push_back(value);
                if (result.second) {
                    uniqueKeys_.push_back(key.value());
                }
            }

            for (const auto &key: uniqueKeys_) {
                auto [keyWriter, valueWriter] = out.add_item();
                keyWriter.copy_from(key);

                const auto &values = keyValuesMap_[key];

                for (const auto &value: values) {
                    valueWriter.push_back(value);
                }
            }
        }

    private:
        melon::F14FastMap<
            exec::GenericView,
            std::vector<std::optional<exec::GenericView> > >
        keyValuesMap_;

        // List of unique keys in the same order as they appear in inputArray.
        // Used to ensure deterministic order of keys in the result. Without ensuring
        // deterministic order of keys, the results of expressions like
        // map_keys(multimap_from_entries(...)) will be non-deterministic and trigger
        // Fuzzer failures. F14FastMap in debug build returns keys in
        // non-deterministic order (on purpose).
        std::vector<exec::GenericView> uniqueKeys_;
    };
} // namespace kumo::pollux::functions
